// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/functions/prestosql/split_to_map.h>

namespace kumo::pollux::functions {
    namespace {
        core::CallTypedExprPtr asSplitToMapCall(
            const std::string &prefix,
            const core::TypedExprPtr &expr) {
            if (auto call = std::dynamic_pointer_cast<const core::CallTypedExpr>(expr)) {
                if (call->name() == prefix + "split_to_map") {
                    return call;
                }
            }
            return nullptr;
        }
    } // namespace

    core::TypedExprPtr rewriteSplitToMapCall(
        const std::string &prefix,
        const core::TypedExprPtr &expr) {
        const auto call = asSplitToMapCall(prefix, expr);
        if (call == nullptr || call->inputs().size() != 4) {
            return nullptr;
        }

        const auto *lambda =
                dynamic_cast<const core::LambdaTypedExpr *>(call->inputs()[3].get());
        POLLUX_CHECK_NOT_NULL(lambda);

        const auto v1 = lambda->signature()->nameOf(1);
        const auto v2 = lambda->signature()->nameOf(2);

        const auto &body = lambda->body();
        std::optional<bool> keepFirst;
        if (auto field =
                std::dynamic_pointer_cast<const core::FieldAccessTypedExpr>(body)) {
            if (field->isInputColumn()) {
                if (field->name() == v1) {
                    keepFirst = true;
                } else if (field->name() == v2) {
                    keepFirst = false;
                }
            }
        }

        if (!keepFirst.has_value()) {
            static const std::string kNotSupported =
                    "split_to_map with arbitrary lambda is not supported: {}";
            POLLUX_USER_FAIL(kNotSupported, lambda->toString());
        }

        return std::make_shared<core::CallTypedExpr>(
            call->type(),
            std::vector<core::TypedExprPtr>{
                call->inputs()[0],
                call->inputs()[1],
                call->inputs()[2],
                std::make_shared<core::ConstantTypedExpr>(
                    BOOLEAN(), keepFirst.value())
            },
            "$internal$split_to_map");
    }
} // namespace kumo::pollux::functions
